package com.yc.web.controller;

import com.yc.config.AiConfig;
import com.yc.config.PersistentMemoryStore;
import com.yc.web.controller.model.ResponseResult;
import dev.langchain4j.data.message.*;
import dev.langchain4j.service.TokenStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.*;
import java.util.stream.Collectors;

@RestController
@RequestMapping("api/ai")
public class AiController {
    @Autowired
    public AiConfig.AiAssistant assistant;
    @Autowired
    private PersistentMemoryStore memoryStore;

    // 删除三角色模式的聊天记录
    @DeleteMapping("/triple-role/history")
    public ResponseResult deleteTripleRoleHistory(
            @RequestParam("memoryId") String memoryId) {
        try {
//            String baseMemoryId = memoryId.split("_")[0];
            memoryStore.deleteTripleRoleData(memoryId);
            return ResponseResult.ok("删除三角色聊天记录成功");
        } catch (Exception e) {
            return ResponseResult.error("删除失败: " + e.getMessage());
        }
    }

    // 删除角色互换模式的聊天记录
    @DeleteMapping("/reverse-role/history")
    public ResponseResult deleteReverseRoleHistory(
            @RequestParam("memoryId") String memoryId) {
        try {
//            String baseMemoryId = memoryId.split("_")[0];
            memoryStore.deletePatientRoleData(memoryId);
            return ResponseResult.ok("删除角色互换聊天记录成功");
        } catch (Exception e) {
            return ResponseResult.error("删除失败: " + e.getMessage());
        }
    }

    // 获取聊天历史数据（按用户ID获取所有AI角色的历史记录）
    @GetMapping("/chat/history")
    public ResponseResult getChatHistory(@RequestParam("userId") String userId) {
        try {
            Set<String> allKeys = memoryStore.getKeysByPattern(userId + "_*");
            Set<String> filteredKeys = new HashSet<>();

            // 定义需要保留的关键词
            Set<String> targetKeywords = Set.of("listener", "coach", "mentor");

            for (String key : allKeys) {
                // 分割键，获取中间部分（假设键的格式是 userId_xxx_patient）
                String[] parts = key.split("_");
                if (parts.length >= 3) {
                    String lastPart = parts[2]; // 最后的部分是我们要判断的内容
                    if (targetKeywords.contains(lastPart)) {
                        filteredKeys.add(key);
                    }
                }
            }

            List<Map<String, Object>> allMessages = new ArrayList<>();

            for (String key : filteredKeys) {
                String[] parts = key.split("_");
                if (parts.length < 3) continue;
                String timestamp = parts[1];
                String role = parts[2];
                List<ChatMessage> messages = memoryStore.getMessages(key);

                for (int i = 0; i < messages.size(); i++) {
                    ChatMessage msg = messages.get(i);

                    if(msg.type() == ChatMessageType.SYSTEM){
                        continue;
                    }

                    Map<String, Object> messageObj = new HashMap<>();

                    // 根据消息类型处理内容
                    String content = extractMessageContent(msg);
                    String sender = determineSender(msg.type());

                    messageObj.put("role", role);
                    messageObj.put("text", content);
//                    messageObj.put("role", sender.equals("USER") ? "user" : role);
                    messageObj.put("type", msg.type().toString());
                    messageObj.put("timestamp", timestamp);

                    allMessages.add(messageObj);
                }
            }

//            allMessages.sort(Comparator.comparingLong(m -> (Long) m.get("timestamp")));
            return ResponseResult.ok("获取历史数据成功").setData(allMessages);

        } catch (Exception e) {
            return ResponseResult.error("获取历史数据失败: " + e.getMessage()).setData(new ArrayList<>());
        }
    }

    // 辅助方法：提取消息内容
    private String extractMessageContent(ChatMessage msg) {
        switch (msg.type()) {
            case SYSTEM:
                return ((SystemMessage) msg).text();
            case USER:
                UserMessage userMessage = (UserMessage) msg;
                // 优先使用 singleText() 方法（当只有单个文本内容时）
                if (userMessage.hasSingleText()) {
                    return userMessage.singleText();
                }
                // 多内容时拼接所有文本内容
                return userMessage.contents().stream()
                        .filter(content -> content instanceof TextContent)
                        .map(content -> ((TextContent) content).text())
                        .collect(Collectors.joining("\n"));
            case AI:
                return ((AiMessage) msg).text();
            default:
                return msg.toString();
        }
    }

    // 辅助方法：确定消息发送者
    private String determineSender(ChatMessageType type) {
        return type == ChatMessageType.USER ? "USER" : "AI";
    }

    // 获取角色互换历史数据
    @GetMapping("/reverse-role/history")
    public ResponseResult getReverseRoleHistory(@RequestParam("userId") String userId) {
        try {
            Set<String> keys = memoryStore.getKeysByPattern(userId + "_*_patient");
            List<Map<String, Object>> allMessages = new ArrayList<>();

            for (String key : keys) {
                String[] parts = key.split("_");
                if (parts.length < 3) continue;
                String timestamp = parts[1];
                String role = parts[2];
                List<ChatMessage> messages = memoryStore.getMessages(key);

                for (int i = 0; i < messages.size(); i++) {
                    ChatMessage msg = messages.get(i);

                    if(msg.type() == ChatMessageType.SYSTEM){
                        continue;
                    }

                    Map<String, Object> messageObj = new HashMap<>();

                    // 根据消息类型处理内容
                    String content = extractMessageContent(msg);
                    String sender = determineSender(msg.type());

//                    messageObj.put("id", role + "_" + i);
                    messageObj.put("text", content);
//                    messageObj.put("role", sender.equals("USER") ? "user" : role);
                    messageObj.put("type", msg.type().toString());
                    messageObj.put("timestamp", timestamp);

                    allMessages.add(messageObj);
                }
            }

//            allMessages.sort(Comparator.comparingLong(m -> (Long) m.get("timestamp")));
            return ResponseResult.ok("获取历史数据成功").setData(allMessages);

        } catch (Exception e) {
            return ResponseResult.error("获取历史数据失败: " + e.getMessage()).setData(new ArrayList<>());
        }
    }

    @GetMapping(value = "/coach/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> assistantStreamingChat(
            @RequestParam("memoryId") String memoryId,
            @RequestParam("question") String question) {

        // 为每个角色生成不同的memoryId
        String listenerMemoryId = memoryId + "_listener";
        String coachMemoryId = memoryId + "_coach";
        String mentorMemoryId = memoryId + "_mentor";

        // 并行调用3个角色的方法
        TokenStream listenerStream = assistant.chatStreamListener(listenerMemoryId, question);
        TokenStream coachStream = assistant.chatStreamTrainer(coachMemoryId, question);
        TokenStream mentorStream = assistant.chatStreamMentalTutor(mentorMemoryId, question);

        // 转换为Flux并合并
        Flux<String> listenerFlux = streamToFlux(listenerStream, "listener");
        Flux<String> coachFlux = streamToFlux(coachStream, "coach");
        Flux<String> mentorFlux = streamToFlux(mentorStream, "mentor");

        return Flux.merge(listenerFlux, coachFlux, mentorFlux);
    }

    // 角色互换模式接口
    @GetMapping(value = "/reverse-role/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> reverseRoleStreamingChat(
            @RequestParam("memoryId") String memoryId,
            @RequestParam("question") String question) {

        // 角色互换使用独立的memoryId
        String patientMemoryId = memoryId + "_patient";

        TokenStream patientStream = assistant.chatStreamPatient(patientMemoryId, question);
        Flux<String> patientFlux = streamToFlux(patientStream, "ai");

        return patientFlux;
    }

    private Flux<String> streamToFlux(TokenStream stream, String role) {
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(response -> {
                String json = String.format("{\"role\":\"%s\",\"text\":%s}", role, toJsonString(response));
                fluxSink.next(json);
            }).onCompleteResponse(response -> {
                String doneJson = String.format("{\"role\":\"%s\",\"done\":true}", role);
                fluxSink.next(doneJson);
                fluxSink.complete();
            }).onError(error -> {
                fluxSink.error(error);
            }).start();
        }, FluxSink.OverflowStrategy.BUFFER);
    }

    private String toJsonString(String s) {
        return "\"" + s.replace("\"", "\\\"").replace("\n", "\\n") + "\"";
    }
}